import torch
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import tkinter as tk
from tkinter import ttk, colorchooser
from PIL import Image  # 提前导入PIL，避免save时临时导入


class PyTorchDrawingBoard:
    def __init__(self, root):
        self.root = root
        self.root.title("PyTorch 基础绘图画板")
        self.root.geometry("800x600")

        # 1. 初始化绘图参数
        self.canvas_size = (600, 400)  # 画布尺寸
        self.brush_size = 5  # 画笔粗细
        self.brush_color = "#000000"  # 初始颜色（黑色）
        self.drawing = False  # 绘图状态标记

        # 2. 创建PyTorch张量作为画布（3通道RGB，值范围[0,1]，初始为白色）
        self.canvas_tensor = torch.ones((*self.canvas_size[::-1], 3))  # (H, W, 3)

        # 3. 构建Matplotlib图形
        self.fig, self.ax = plt.subplots(figsize=(self.canvas_size[0]/100, self.canvas_size[1]/100))
        self.ax.set_xlim(0, self.canvas_size[0])
        self.ax.set_ylim(0, self.canvas_size[1])
        self.ax.invert_yaxis()  # 反转y轴（符合屏幕坐标）
        self.ax.axis("off")  # 隐藏坐标轴

        # 4. 绑定鼠标事件
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
        self.canvas.draw()
        self.canvas.get_tk_widget().pack(side=tk.LEFT, padx=10, pady=10)
        self.canvas.mpl_connect("button_press_event", self.on_press)
        self.canvas.mpl_connect("motion_notify_event", self.on_drag)
        self.canvas.mpl_connect("button_release_event", self.on_release)

        # 5. 右侧控制面板
        self.control_frame = ttk.Frame(self.root, padding=10)
        self.control_frame.pack(side=tk.RIGHT, fill=tk.Y)

        # 画笔粗细调节
        ttk.Label(self.control_frame, text="画笔粗细:").pack(pady=5)
        self.brush_size_var = tk.IntVar(value=self.brush_size)
        brush_slider = ttk.Scale(
            self.control_frame, from_=1, to=20, variable=self.brush_size_var,
            orient=tk.HORIZONTAL, command=self.update_brush_size
        )
        brush_slider.pack(pady=5, fill=tk.X)

        # 颜色选择按钮
        self.color_btn = ttk.Button(
            self.control_frame, text="选择颜色", command=self.choose_color
        )
        self.color_btn.pack(pady=5, fill=tk.X)
        self.color_label = ttk.Label(self.control_frame, text=f"当前颜色: {self.brush_color}")
        self.color_label.pack(pady=2)

        # 清空画布按钮
        clear_btn = ttk.Button(
            self.control_frame, text="清空画布", command=self.clear_canvas
        )
        clear_btn.pack(pady=5, fill=tk.X)

        # 保存画布按钮
        save_btn = ttk.Button(
            self.control_frame, text="保存画布", command=self.save_canvas
        )
        save_btn.pack(pady=5, fill=tk.X)

    # 鼠标按下事件（开始绘图）
    def on_press(self, event):
        if event.xdata is not None and event.ydata is not None:
            self.drawing = True
            self.last_x, self.last_y = event.xdata, event.ydata

    # 鼠标拖动事件（绘制线条）
    def on_drag(self, event):
        if self.drawing and event.xdata is not None and event.ydata is not None:
            # 计算绘制区域（基于画笔粗细）
            x1, y1 = self.last_x, self.last_y
            x2, y2 = event.xdata, event.ydata

            # 生成线段上的所有点（用numpy计算，再转为PyTorch张量）
            num_points = int(max(abs(x2-x1), abs(y2-y1)) * 2)
            if num_points == 0:
                num_points = 1
            x_points = np.linspace(x1, x2, num_points)
            y_points = np.linspace(y1, y2, num_points)

            # 将坐标转为整数（对应张量索引）
            x_indices = torch.clamp(torch.tensor(x_points).int(), 0, self.canvas_size[0]-1)
            y_indices = torch.clamp(torch.tensor(y_points).int(), 0, self.canvas_size[1]-1)

            # 绘制圆形画笔痕迹（基于画笔粗细）
            for x, y in zip(x_indices, y_indices):
                # 生成圆形区域的坐标
                y_range = torch.arange(max(0, y - self.brush_size//2), min(self.canvas_size[1], y + self.brush_size//2 + 1))
                x_range = torch.arange(max(0, x - self.brush_size//2), min(self.canvas_size[0], x + self.brush_size//2 + 1))
                yy, xx = torch.meshgrid(y_range, x_range, indexing="ij")

                # 计算距离，小于半径的点设为画笔颜色
                distance = torch.sqrt((xx - x)**2 + (yy - y)**2)
                mask = distance <= self.brush_size / 2

                # 将颜色十六进制转为RGB张量（范围[0,1]）
                rgb = torch.tensor([int(self.brush_color[i:i+2], 16)/255.0 for i in (1,3,5)])
                self.canvas_tensor[yy[mask], xx[mask]] = rgb

            # 更新画布
            self.update_canvas()
            self.last_x, self.last_y = x2, y2

    # 鼠标释放事件（结束绘图）
    def on_release(self, event):  # 修复：缩进为类内方法（1层）
        self.drawing = False       # 函数体缩进（2层）

    # 更新画笔粗细
    def update_brush_size(self, value):  # 修复：缩进为类内方法
        self.brush_size = int(float(value))

    # 选择颜色
    def choose_color(self):  # 修复：缩进为类内方法
        color = colorchooser.askcolor(title="选择画笔颜色")[1]
        if color:
            self.brush_color = color
            self.color_label.config(text=f"当前颜色: {self.brush_color}")

    # 清空画布（重置为白色张量）
    def clear_canvas(self):  # 修复：缩进为类内方法
        self.canvas_tensor = torch.ones((*self.canvas_size[::-1], 3))
        self.update_canvas()

    # 保存画布（转为PIL图像保存）
    def save_canvas(self):  # 修复：缩进为类内方法
        # 张量转为numpy数组，范围[0,255]，uint8格式
        img_np = (self.canvas_tensor.numpy() * 255).astype(np.uint8)
        img = Image.fromarray(img_np)
        img.save("drawing_result.png")
        print("画布已保存为 drawing_result.png")

    # 更新Matplotlib显示
    def update_canvas(self):  # 修复：缩进为类内方法
        self.ax.clear()
        self.ax.imshow(self.canvas_tensor)
        self.ax.axis("off")
        self.canvas.draw()


if __name__ == "__main__":
    root = tk.Tk()
    app = PyTorchDrawingBoard(root)
    root.mainloop()